from django.db.models import Q
from django_filters import rest_framework as df

from common.models import Company


class CompanyFilterSet(df.FilterSet):
    """自定义过滤类"""
    icate = df.NumberFilter(field_name='industry_category')
    fstage = df.NumberFilter(field_name='financing_stage')
    ssize = df.NumberFilter(field_name='staff_size')
    name = df.CharFilter(method='filter_by_name')

    @staticmethod
    def filter_by_name(queryset, name, value):
        return queryset.filter(
            Q(short_name__contains=value) | Q(full_name__contains=value)
        )

    class Meta:
        model = Company
        fields = ('icate', 'fstage', 'ssize', 'name')
